from itertools import combinations
from collections import Counter
from mult_align import *
from myseq import MySeq
from align import SubstMatrix
from bio_sm import SubstMatrix_Bio


def multiple_align_case1():
    s1 = MySeq("ACATATCAT")
    s2 = MySeq("AACAGATCT")
    s3 = MySeq("AGATATTAG")
    s4 = MySeq("GCATCGATT")
    sm = SubstMatrix()
    sm.create_submat(1,-1,MySeq.ALPHABET["DNA"])
    aligner = PairwiseAlignment(sm,-1)
    multi_aligner = MultipleAlignment([s1,s2,s3,s4],aligner)
    print(multi_aligner.align_consensus())

class MultipleAlignment_ex(MultipleAlignment):

    def score_pos(self,chars:list):
        aliner = self.align_pars
        score = aliner.gap*chars.count('-')
        for c1,c2 in combinations(chars,2):
            if c1 == '-' or c2 == '-':
                continue
            else:
                score += aliner.sm[c1,c2]
        return score
    
    def score_sum_pairs(self,align:MyAlign) -> int:
        self.alignpars = 0
        for index in range(len(align)):
            chars = align.column(index)
            self.alignpars += self.score_pos(chars)
        return self.alignpars
    
    def high_quality_colum(self,align:MyAlign) -> list[bool]:
        high_quality = []
        for index in range(len(align)):
            chars = align.column(index)
            high_quality.append(self.score_pos(chars) > 0)
        return high_quality

def test_SP():
    align = MyAlign(['-RADNS','ARCD-A','AR-D-A'])
    sm = SubstMatrix_Bio()
    aligner = PairwiseAlignment(sm,-8)
    multi_aligner = MultipleAlignment_ex([],aligner)
    print(multi_aligner.score_sum_pairs(align))

def multiple_align_case2():
    print("a part:")
    s1 = MySeq('APSC')
    s2 = MySeq('TAPT')
    sm = SubstMatrix_Bio()
    aligner = PairwiseAlignment(sm,-4)
    aligner.needleman_Wunsch(s1,s2)
    print(aligner.recover_align())

    print("b part:")
    s1 = MySeq('APSC')
    s2 = MySeq('TAPT')
    s3 = MySeq('TATS')
    multi_aligner = MultipleAlignment_ex([s1,s2,s3],aligner)
    align = multi_aligner.align_consensus()
    print(align)

    print('c part:')
    print(f"SP: {multi_aligner.score_sum_pairs(align)}")

class MyAlign_ex(MyAlign):
    def polar_rich(self):
        polar_base = ('R','H','K')
        result = []
        for index in range(len(self)):
            chars = self.column(index)
            counter = Counter(chars)
            num_polar = sum((counter[b] for b in polar_base))
            if num_polar*2 >= self.num_seqs():
                result.append(index)
        return result
    
    def symbol_str(self) -> str:
        result = ''
        for index in range(len(self)):
            chars = self.column(index)
            counter = Counter(chars)
            max_base,max_count = counter.most_common(1)[0]
            if len(counter) == 1:
                result += '*'
            elif max_base!='-' and max_count*2 >= self.num_seqs():
                result += ':'
            elif not counter.get('-',False):
                result += '.'
            else:
                result += ' '
        return result

def test_polar_rich():
    align = MyAlign_ex(['-RADNS','ARCDHA','AR-DKA'])
    print(align.polar_rich())

def test_symbol_str():
    align = MyAlign_ex(['-RADNS','ARCD-A','ARND-A'])
    print(align.symbol_str())

def test_high_quality():
    align = MyAlign(['-RADNS','ARCD-A','ARND-A'])
    sm = SubstMatrix_Bio()
    aligner = PairwiseAlignment(sm,-8)
    multi_aligner = MultipleAlignment_ex([],aligner)
    print(multi_aligner.high_quality_colum(align))


if __name__ == "__main__":
    print("### ex1 ###")
    multiple_align_case1()
    print('### ex2 ###')
    test_SP()
    print('### ex3 ###')
    multiple_align_case2()
    print('### ex4 ###')
    test_polar_rich()
    print('### ex5 ###')
    test_symbol_str()
    print('### ex6 ###')
    test_high_quality()